#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ROS 2 bag 延时分析器：计算 Δ1, Δ2, Δ3
- Δ1 = cmd.stamp - wrench.stamp
- Δ2 = joint_state(首次变化).stamp - cmd.stamp
- Δ3 = joint_state(首次变化).stamp - wrench.stamp

使用方法：
  1) 先在终端 source ROS 2 环境，比如：
       source /opt/ros/humble/setup.bash
  2) 运行：
       python3 analyze_latency_ros2bag.py -b /path/to/rosbag2_2025_09_29-13_15_03 \
         --wrench /wrench_comp \
         --cmd /force_servo_control_node/delta_twist_cmds \
         --joint /joint_states
  3) 如果你的 cmd 话题或力话题名不同，用 --cmd /xxx --wrench /yyy 指定
"""

import argparse
import math
from typing import Dict, Optional, List

# 需要 ROS 2 Python 接口
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions
from rclpy.serialization import deserialize_message
from rosidl_runtime_py.utilities import get_message


def builtin_time_to_sec(t) -> float:
    """builtin_interfaces/Time -> float seconds"""
    return float(t.sec) + float(t.nanosec) * 1e-9


def ns_to_sec(n: int) -> float:
    return float(n) * 1e-9


def main():
    parser = argparse.ArgumentParser(description="Compute Δ1/Δ2/Δ3 from ROS 2 bag")
    parser.add_argument("-b", "--bag", required=True, help="rosbag 目录（不是 _0.db3 文件路径！）")
    parser.add_argument("--wrench", default="/wrench_comp", help="力话题（WrenchStamped）")
    parser.add_argument("--cmd", default="/force_servo_control_node/delta_twist_cmds", help="命令话题（TwistStamped）")
    parser.add_argument("--joint", default="/joint_states", help="关节状态（sensor_msgs/JointState）")
    parser.add_argument("--eps_cmd", type=float, default=1e-3, help="判定cmd从0到非0的阈值（m/s）")
    parser.add_argument("--vel_thresh", type=float, default=1e-3, help="判定关节“开始运动”的速度阈值（rad/s）")
    parser.add_argument("--pos_deriv_thresh", type=float, default=1e-3, help="若JointState.velocity无效，则用位置微分判定的阈值（rad/s）")
    args = parser.parse_args()

    storage_options = StorageOptions(uri=args.bag, storage_id="sqlite3")
    converter_options = ConverterOptions(input_serialization_format="cdr", output_serialization_format="cdr")
    reader = SequentialReader()
    reader.open(storage_options, converter_options)

    # 获取所有话题类型
    topic_types: Dict[str, str] = {}
    for t in reader.get_all_topics_and_types():
        topic_types[t.name] = t.type

    needed = [args.wrench, args.cmd, args.joint]
    for n in needed:
        if n not in topic_types:
            raise RuntimeError(f"包内未找到话题 {n}，可用话题有：{list(topic_types.keys())}")

    # 准备反序列化
    msg_types = {name: get_message(topic_types[name]) for name in needed}

    # 逐条读取
    last_wrench_stamp_ros: Optional[float] = None  # wrench 的 header.stamp（秒）
    found_cmd = False
    cmd_stamp_ros: Optional[float] = None
    js_prev_pos: Optional[List[float]] = None
    js_prev_stamp: Optional[float] = None
    js_change_stamp_ros: Optional[float] = None

    # 统计计数
    cnt_wrench = cnt_cmd = cnt_js = 0

    # 为了检测“首次 cmd 非零”，记录上一帧 cmd.vx
    prev_cmd_vx: Optional[float] = 0.0

    while reader.has_next():
        topic, data, t = reader.read_next()  # t 是接收时间（ns），本脚本使用 header.stamp 做计算
        if topic not in msg_types:
            continue
        msg = deserialize_message(data, msg_types[topic])

        if topic == args.wrench:
            cnt_wrench += 1
            # 取 header.stamp（如果没有，就用 t）
            try:
                stamp_ros = builtin_time_to_sec(msg.header.stamp)
            except Exception:
                stamp_ros = ns_to_sec(t)
            last_wrench_stamp_ros = stamp_ros

        elif topic == args.cmd and not found_cmd:
            cnt_cmd += 1
            try:
                stamp_ros = builtin_time_to_sec(msg.header.stamp)
            except Exception:
                stamp_ros = ns_to_sec(t)

            try:
                vx = float(msg.twist.linear.x)
            except Exception:
                vx = 0.0

            # 检测从“近零”->“超过阈值”的首次跃迁
            if prev_cmd_vx is None:
                prev_cmd_vx = vx
            if (abs(vx) > args.eps_cmd) and (abs(prev_cmd_vx or 0.0) <= args.eps_cmd):
                found_cmd = True
                cmd_stamp_ros = stamp_ros
            prev_cmd_vx = vx

        elif topic == args.joint and (found_cmd and js_change_stamp_ros is None):
            cnt_js += 1
            # JointState 可能有 velocity；如果没有就用位置微分
            try:
                js_stamp = builtin_time_to_sec(msg.header.stamp)
            except Exception:
                js_stamp = ns_to_sec(t)

            has_velocity = hasattr(msg, "velocity") and isinstance(msg.velocity, (list, tuple)) and len(msg.velocity) == len(msg.position)

            triggered = False

            if has_velocity:
                for v in msg.velocity:
                    try:
                        if abs(float(v)) >= args.vel_thresh:
                            triggered = True
                            break
                    except Exception:
                        pass
            else:
                # 位置微分
                if js_prev_pos is not None and js_prev_stamp is not None:
                    dt = js_stamp - js_prev_stamp
                    if dt > 0.0:
                        for p, pp in zip(msg.position, js_prev_pos):
                            try:
                                vel = (float(p) - float(pp)) / dt
                                if abs(vel) >= args.pos_deriv_thresh:
                                    triggered = True
                                    break
                            except Exception:
                                pass

            js_prev_pos = list(msg.position) if hasattr(msg, "position") else None
            js_prev_stamp = js_stamp

            # 只在 cmd 之后触发
            if triggered and js_stamp >= (cmd_stamp_ros or js_stamp):
                js_change_stamp_ros = js_stamp

        # 如果全拿到了可提前退出
        if found_cmd and (js_change_stamp_ros is not None) and (last_wrench_stamp_ros is not None):
            break

    print("=== 统计 ===")
    print(f"wrench msgs: {cnt_wrench}, cmd msgs: {cnt_cmd}, joint msgs: {cnt_js}")
    if last_wrench_stamp_ros is None:
        print("未捕获到任何 wrench 帧。")
        return
    if not found_cmd or cmd_stamp_ros is None:
        print("未检测到 cmd 的“首次非零跃迁”。（可调小 --eps_cmd 或检查你的命令话题）")
        return
    if js_change_stamp_ros is None:
        print("未检测到 cmd 之后的关节“首次明显变化”。（可调低 --vel_thresh 或 --pos_deriv_thresh）")
        return

    # 计算 Δ1 Δ2 Δ3（毫秒）
    d1 = (cmd_stamp_ros - last_wrench_stamp_ros) * 1000.0
    d2 = (js_change_stamp_ros - cmd_stamp_ros) * 1000.0
    d3 = (js_change_stamp_ros - last_wrench_stamp_ros) * 1000.0

    print("\n=== 结果（毫秒 ms）===")
    print(f"Δ1 = cmd.stamp - wrench.stamp = {d1:.2f} ms")
    print(f"Δ2 = joint_change - cmd.stamp   = {d2:.2f} ms")
    print(f"Δ3 = joint_change - wrench.stamp= {d3:.2f} ms")

    print("\n时间戳（秒）参考：")
    print(f"  wrench_stamp       = {last_wrench_stamp_ros:.6f}")
    print(f"  cmd_first_nonzero  = {cmd_stamp_ros:.6f}")
    print(f"  joint_first_change = {js_change_stamp_ros:.6f}")


if __name__ == "__main__":
    main()
